#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @time: 2023/11/15 14:58
# @file: get_tick_data.py
# @author: tyshixi08


import os
import rqdatac
import pandas as pd
import datetime
from datetime import date


def bond_convert(code):
    """
    从米筐合约代码转换至交易所合约代码(米筐仅支持转换A股、期货和期权)
    """
    if code[7:] == 'XSHG':
        code = code[:7] + 'SH'
    elif code[7:] == 'XSHE':
        code = code[:7] + 'SZ'
    return code


def get_bond_data(code, start, end):
    """
    获取可转债5分钟历史行情数据
    """
    data = rqdatac.get_price(code, start_date=start, end_date=end, frequency='1m', fields=None,
                             adjust_type='none', skip_suspended=False, market='cn', expect_df=True, time_slice=None)
    try:
        data = data.reset_index()
        data['Date'] = data['datetime'].apply(lambda x: x.strftime('%Y-%m-%d'))
        data = data[['order_book_id', 'datetime', 'low', 'high', 'close','open', 'total_turnover', 'volume', 'Date']]
        data = data.rename(columns={'order_book_id': 'c_bondCode', 'open': 'bond_open', 'close': 'bond_close'})
        data['c_bondCode'] = data['c_bondCode'].apply(bond_convert)
    except:
        print('数据已更到最新')
    return data


def get_stock_data(code, start, end):
    """
    获取股票5分钟历史行情数据
    """
    data = rqdatac.get_price(code, start_date=start, end_date=end, frequency='1m', fields=None,
                             adjust_type='none', skip_suspended=False, market='cn', expect_df=True, time_slice=None)
    try:
        data = data.reset_index()
        data['Date'] = data['datetime'].apply(lambda x: x.strftime('%Y-%m-%d'))
        data = data[['order_book_id', 'datetime', 'low', 'high', 'open', 'close', 'Date']]
        data = data.rename(columns={'order_book_id': 'c_stockCode', 'open': 'stock_open', 'close': 'stock_close'})
        data['c_stockCode'] = data['c_stockCode'].apply(bond_convert)
    except:
        print('数据已更到最新')
    return data


def get_open_data(code, start, end, name):
    """
    获取09:25集合竞价数据
    """
    data = rqdatac.get_open_auction_info(code, start, end)
    try:
        data = data.reset_index()
        data['Date'] = data['datetime'].apply(lambda x: x.strftime('%Y-%m-%d'))
        data['datetime'] = data['Date'] + ' 09:25:00'
        data = data.rename(columns={'order_book_id': 'c_{}Code'.format(name), 'open': '{}_open'.format(name),
                                    'last': '{}_close'.format(name)})
        data['c_{}Code'.format(name)] = data['c_{}Code'.format(name)].apply(bond_convert)
    except:
        print('数据已更到最新')
    return data


def convert_data(port, field): # 将1分钟的数据转成5分钟的，并且保留9：31
    port['tail'] = [datetime.date.strftime(x,'%H:%M:%S') for x in pd.to_datetime(port['datetime'])]
    if field == 'bond':
        port['{}_open_nextmin'.format(field)] = port.groupby(['c_{}Code'.format(field)])['{}_open'.format(field)].shift(-1)
        port['{}_close_nextmin'.format(field)] = port.groupby(['c_{}Code'.format(field)])['{}_close'.format(field)].shift(-1)
    tlst = []    
    tail = port['tail'].unique()
    for i in range(len(tail)):
        if (i+1) % 5 == 0:
            tlst.append(tail[i])
    tlst.insert(0,'09:31:00')
    port = port[port['tail'].isin(tlst)]
    return port


def get_data():
    end_date = date.today() - datetime.timedelta(days=1)
    end_date = end_date.strftime('%Y-%m-%d')
    rqdatac.init('license', 'AcBHy5_JJ6wjZdu7Q-ey7dX-J3BmyEC_KblY2Q_hBeOuoBaeBbgXTNSe6XZvqKVESbyUf7vMpLLGuO_aqyb3w9fWGI7q4wdClE6cMp_Z3N4PqqTHJ0nr3CIuXtk-5XzSD1p7NTdNcrAfZlRVpMMtY_PDC9FYuXNmC_EnuQg4H-A=fGk9EhHcK3xN189iXYSWLyiMdGUeXXlVZqr2MxhBypSHxQYnIIyxyM8BR8oNnVUdWhKx-ZrFRIjSONd7uYpOvpcBab92P60iAR_JopX61emtrvsY1xG_uCfYhDPBdDSJKaniJhTPuoBIU4JZun8-8fMIxzx7lnwBm2kAUOA_Mpg=')

    # 可转债基础信息
    print('instrument')
    instrument = rqdatac.all_instruments(type='Convertible', market='cn', date=None)
    instrument.to_csv('./data/instrument.csv', index=False, encoding='utf_8_sig')

    # 可转债和股票代码
    bond_code = instrument['order_book_id'].unique()
    stock_code0 = instrument['stock_code'].unique()
    stock_code = [code if not pd.isnull(code) else None for code in stock_code0]
    while None in stock_code:
        stock_code.remove(None)

    # 可转债5分钟历史行情数据 
    print('bond_data')
    bond_data = pd.read_pickle(os.path.join('data', 'bond_data.pkl'))
    start_date = pd.to_datetime(list(bond_data.Date.unique())[-1]).strftime('%Y-%m-%d')
    bond_data_supple = get_bond_data(bond_code, start_date, end_date)
    try:
        bond_data_supple = convert_data(bond_data_supple, 'bond')
        bond_data = pd.concat([bond_data[bond_data.Date <= list(bond_data.Date.unique())[-2]], bond_data_supple],axis=0)
    except:
        pass
    print(len(bond_data))
    bond_data = bond_data.drop_duplicates(['c_bondCode','datetime'], keep='last')
    print(len(bond_data))
    bond_data = bond_data.sort_values(['datetime','c_bondCode'])
    last_dt = pd.to_datetime(list(bond_data.datetime.unique())[-1])
    print(last_dt)
    if last_dt.strftime('%H:%M:%S') == '15:00:00':
        bond_data.to_pickle(os.path.join('data', 'bond_data.pkl'))
    else:
        print('当日数据未完全更新')
    # 补充首次提取代码：
    # bond_data = get_bond_data(bond_code, start_date, end_date)
    # bond_data = convert_data(bond_data, 'bond')
    # bond_data = bond_data.drop_duplicates(['c_bondCode','datetime'])
    # bond_data = bond_data.sort_values(['datetime','c_bondCode'])
    # bond_data.to_pickle(os.path.join('data', 'bond_data.pkl'))
    
    # 股票5分钟历史行情数据
    print('stock_data')
    stock_data = pd.read_pickle(os.path.join('data', 'stock_data.pkl'))
    start_date = pd.to_datetime(stock_data.Date.unique()[-1])
    start_date = start_date.strftime('%Y-%m-%d')
    stock_data_supple = get_stock_data(stock_code, start_date, end_date)
    try:
        stock_data_supple = convert_data(stock_data_supple, 'stock')
        stock_data = pd.concat([stock_data[stock_data.Date <= list(stock_data.Date.unique())[-2]], stock_data_supple],axis=0)
    except:
        pass
    print(len(stock_data))
    stock_data = stock_data.drop_duplicates(['c_stockCode','datetime'])
    print(len(stock_data))
    stock_data = stock_data.sort_values(['datetime','c_stockCode'])
    last_dt = pd.to_datetime(list(stock_data.datetime.unique())[-1])
    print(last_dt)
    if last_dt.strftime('%H:%M:%S') == '15:00:00':
        stock_data.to_pickle(os.path.join('data', 'stock_data.pkl'))
    else:
        print('当日数据未完全更新')
    # 补充首次提取代码：
    # stock_data = get_stock_data(stock_code, start_date, end_date)
    # stock_data = convert_data(stock_data, 'stock')
    # stock_data = stock_data.drop_duplicates(['c_stockCode','datetime'])
    # stock_data = stock_data.sort_values(['datetime','c_stockCode'])
    # stock_data.to_pickle(os.path.join('data', 'stock_data.pkl'))

    # 可转债集合竞价数据
    print('bond_data_open')
    bond_data_open = pd.read_pickle(os.path.join('data', 'bond_data_open.pkl'))
    start_date = pd.to_datetime(bond_data_open.Date.unique()[-1])
    start_date = start_date.strftime('%Y-%m-%d')
    bond_data_open_supple = get_open_data(bond_code, start_date, end_date, 'bond')
    try:
        bond_data_open_supple = bond_data_open_supple[['c_bondCode', 'datetime', 'low', 'high', 'bond_open',
                                                       'bond_close', 'total_turnover', 'volume', 'Date']]
        bond_data_open = pd.concat([bond_data_open,  bond_data_open_supple],axis=0)
    except:
        pass
    bond_data_open = bond_data_open.drop_duplicates(['c_bondCode','datetime'])
    bond_data_open = bond_data_open.sort_values(['datetime','c_bondCode'])
    bond_data_open.to_pickle(os.path.join('data', 'bond_data_open.pkl'))
    print(pd.to_datetime(list(bond_data_open.datetime.unique())[-1]))

    # 股票集合竞价数据
    print('stock_data_open')
    stock_data_open = pd.read_pickle(os.path.join('data', 'stock_data_open.pkl'))
    start_date = pd.to_datetime(stock_data_open.Date.unique()[-1])
    start_date = start_date.strftime('%Y-%m-%d')
    stock_data_open_supple = get_open_data(stock_code, start_date, end_date, 'stock')
    try:
        stock_data_open_supple = stock_data_open_supple[['c_stockCode', 'datetime', 'low', 'high', 'stock_open',
                                                         'stock_close', 'Date']]
        stock_data_open = pd.concat([stock_data_open, stock_data_open_supple],axis=0)
    except:
        pass
    stock_data_open = stock_data_open.drop_duplicates(['c_stockCode','datetime'])
    stock_data_open = stock_data_open.sort_values(['datetime','c_stockCode'])
    stock_data_open.to_pickle(os.path.join('data', 'stock_data_open.pkl'))
    print(pd.to_datetime(list(stock_data_open.datetime.unique())[-1]))
    return

if __name__ == '__main__':
    get_data()
